package com.ocean.everyday.d002_findingmkaverage;

import java.util.Deque;
import java.util.LinkedList;
import java.util.TreeMap;

/**
 * https://leetcode.cn/problems/finding-mk-average/
 * 求出MK平均值
 */
public class MKAverage2 {

    private int m;
    private int k;
    private TreeMap<Integer, Integer> s1;
    private TreeMap<Integer, Integer> s2;
    private TreeMap<Integer, Integer> s3;
    private int size1;
    private int size2;
    private int size3;
    private int sum2;
    private Deque<Integer> dequeue;

    public static void main(String[] args) {
        MKAverage2 obj = new MKAverage2(3, 1);
        obj.addElement(3);        // 当前元素为 [3]
        obj.addElement(1);        // 当前元素为 [3,1]
        int i = obj.calculateMKAverage();// 返回 -1 ，因为 m = 3
        System.out.println(i);
        obj.addElement(10);
        i = obj.calculateMKAverage();
        System.out.println(i);

        obj.addElement(5);        // 当前元素为 [3,1,10,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5,5]
        i = obj.calculateMKAverage();
        System.out.println(i);
    }


    public MKAverage2(int m, int k) {
        this.m = m;
        this.k = k;
        this.size1 = 0;
        this.size2 = 0;
        this.size3 = 0;
        this.sum2 = 0;
        this.s1 = new TreeMap<>();
        this.s2 = new TreeMap<>();
        this.s3 = new TreeMap<>();
        this.dequeue = new LinkedList<>();

    }

    public void addElement(int num) {
        // 分两种情况 第一种情况：元素的个数<=m  第二种情况元素的个数>m
        dequeue.offer(num);
        if (dequeue.size() <= m) {
            put(s2, num);
            sum2 += num;
            size2++;
            // 如果刚好有m个 就要分到不同的集合
            if (dequeue.size() == m) {
                while (size1 < k) {
                    Integer firstKey = s2.firstKey();
                    remove(s2, firstKey);
                    sum2 -= firstKey;
                    size2--;
                    put(s1, firstKey);
                    size1++;
                }
                while (size3 < k) {
                    Integer lastKey = s2.lastKey();
                    remove(s2, lastKey);
                    sum2 -= lastKey;
                    size2--;
                    put(s3, lastKey);
                    size3++;
                }
            }
            return;
        }
        // m+1个是判断这个元素放到哪个集合
        if (num < s1.lastKey()) {
            Integer lastKey = s1.lastKey();
            put(s1, num);
            remove(s1, lastKey);
            put(s2, num);
            size2++;
            sum2 += num;
        } else if (num > s3.firstKey()) {
            Integer firstKey = s3.firstKey();
            put(s3, num);
            remove(s3, num);
            put(s2, firstKey);
            size2++;
            sum2 += firstKey;
        } else {
            put(s2, num);
            size2++;
            sum2 += num;
        }
        Integer poll = dequeue.poll();
        // 判断从哪个集合中移除
        if (s1.containsKey(poll)) {
            Integer firstKey = s2.firstKey();
            remove(s1, poll);
            put(s1, firstKey);
            remove(s2, firstKey);
            size2--;
            sum2 -= firstKey;
        } else if (s3.containsKey(poll)) {
            Integer lastKey = s2.lastKey();
            remove(s3, poll);
            put(s3, lastKey);
            remove(s2, lastKey);
            size2--;
            sum2 -= lastKey;
        } else {
            remove(s2, poll);
            size2--;
            sum2 -= poll;
        }

    }

    /**
     * @return
     */
    public int calculateMKAverage() {
        if (dequeue.size() < m) {
            return -1;
        }
        return (int) sum2 / (m - 2 * k);
    }

    private void put(TreeMap<Integer, Integer> treeMap, int num) {
        treeMap.put(num, treeMap.getOrDefault(num, 0) + 1);
    }

    private void remove(TreeMap<Integer, Integer> treeMap, int num) {
        treeMap.put(num, treeMap.get(num) - 1);
        if (treeMap.getOrDefault(num, 0) == 0) {
            treeMap.remove(num);
        }
    }
}
